package com.frp.client;

import com.frp.config.FrpcConfig;
import com.frp.protocol.*;
import com.frp.util.CryptoUtil;
import com.frp.util.IdGenerator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.SocketChannel;
import java.util.Iterator;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * 连接管理器
 * 管理与FRP服务器的控制连接
 */
public class ConnectionManager {
    private static final Logger logger = LoggerFactory.getLogger(ConnectionManager.class);
    private static final int BUFFER_SIZE = 8192;
    private static final int MAX_RETRY_ATTEMPTS = 3;
    private static final int RETRY_DELAY_MS = 5000;
    
    private final FrpcConfig config;
    private SocketChannel controlChannel;
    private Selector selector;
    private final AtomicBoolean running = new AtomicBoolean(false);
    private final ConcurrentHashMap<String, ProxyHandler> proxyHandlers = new ConcurrentHashMap<>();
    private final ByteBuffer readBuffer = ByteBuffer.allocate(BUFFER_SIZE);
    private String runId;
    
    public ConnectionManager(FrpcConfig config) {
        this.config = config;
    }
    
    /**
     * 启动连接管理器
     */
    public void start() throws IOException {
        if (running.compareAndSet(false, true)) {
            logger.info("启动连接管理器...");
            connect();
            startMessageLoop();
        }
    }
    
    /**
     * 停止连接管理器
     */
    public void stop() {
        if (running.compareAndSet(true, false)) {
            logger.info("停止连接管理器...");
            try {
                if (controlChannel != null) {
                    controlChannel.close();
                }
                if (selector != null) {
                    selector.wakeup();
                    selector.close();
                }
            } catch (IOException e) {
                logger.error("关闭连接时出错", e);
            }
        }
    }
    
    /**
     * 连接到FRP服务器
     */
    private void connect() throws IOException {
        int retryCount = 0;
        while (retryCount < MAX_RETRY_ATTEMPTS && running.get()) {
            try {
                selector = Selector.open();
                controlChannel = SocketChannel.open();
                controlChannel.configureBlocking(false);
                
                InetSocketAddress address = new InetSocketAddress(
                    config.getCommon().getServerAddr(), 
                    config.getCommon().getServerPort()
                );
                
                boolean connected = controlChannel.connect(address);
                if (!connected) {
                    while (!controlChannel.finishConnect()) {
                        Thread.sleep(100);
                    }
                }
                
                controlChannel.register(selector, SelectionKey.OP_READ);
                
                // 发送登录消息
                sendLoginMessage();
                
                logger.info("成功连接到FRP服务器: {}:{}", 
                    config.getCommon().getServerAddr(), 
                    config.getCommon().getServerPort());
                return;
                
            } catch (Exception e) {
                retryCount++;
                logger.error("连接失败，重试 {}/{}: {}", 
                    retryCount, MAX_RETRY_ATTEMPTS, e.getMessage());
                
                try {
                    Thread.sleep(RETRY_DELAY_MS);
                } catch (InterruptedException ie) {
                    Thread.currentThread().interrupt();
                    break;
                }
            }
        }
        
        throw new IOException("无法连接到FRP服务器");
    }
    
    /**
     * 发送登录消息
     */
    // 发送登录消息
    private void sendLoginMessage() throws IOException {
        LoginMessage loginMsg = new LoginMessage();
        loginMsg.setVersion("0.1.0");
        loginMsg.setHostname(System.getProperty("os.name"));
        loginMsg.setOs(System.getProperty("os.name"));
        loginMsg.setArch(System.getProperty("os.arch"));
        loginMsg.setUser(System.getProperty("user.name"));
        loginMsg.setTimestamp(System.currentTimeMillis());
        loginMsg.setRunId(IdGenerator.generateRunId());
        loginMsg.setPoolCount(config.getCommon().getPoolCount());
        
        // 生成特权密钥
        String privilegeKey = CryptoUtil.generatePrivilegeKey(
            config.getCommon().getToken(), 
            String.valueOf(loginMsg.getTimestamp())
        );
        loginMsg.setPrivilegeKey(privilegeKey);
        
        sendMessage(loginMsg);
    }
    
    /**
     * 启动消息循环
     */
    private void startMessageLoop() {
        Thread messageLoopThread = new Thread(() -> {
            while (running.get()) {
                try {
                    if (selector.select(1000) > 0) {
                        Set<SelectionKey> selectedKeys = selector.selectedKeys();
                        Iterator<SelectionKey> iter = selectedKeys.iterator();
                        
                        while (iter.hasNext()) {
                            SelectionKey key = iter.next();
                            iter.remove();
                            
                            if (key.isReadable()) {
                                handleRead(key);
                            }
                        }
                    }
                } catch (IOException e) {
                    logger.error("消息循环错误", e);
                    if (running.get()) {
                        try {
                            reconnect();
                        } catch (IOException ex) {
                            logger.error("重连失败", ex);
                        }
                    }
                }
            }
        }, "MessageLoopThread");
        
        messageLoopThread.start();
    }
    
    /**
     * 处理读取事件
     */
    private void handleRead(SelectionKey key) {
        try {
            SocketChannel channel = (SocketChannel) key.channel();
            readBuffer.clear();
            
            int bytesRead = channel.read(readBuffer);
            if (bytesRead == -1) {
                logger.warn("连接关闭");
                reconnect();
                return;
            }
            
            if (bytesRead > 0) {
                readBuffer.flip();
                
                Message message = MessageCodec.decode(readBuffer);
                handleMessage(message);
            }
        } catch (Exception e) {
            logger.error("处理读取事件错误", e);
        }
    }
    
    /**
     * 处理消息
     */
    private void handleMessage(Message message) {
        if (message == null) {
            return;
        }
        
        logger.debug("收到消息: {}", message.getType());
        
        switch (message.getType()) {
            case LoginResp:
                handleLoginResponse((LoginRespMessage) message);
                break;
            case NewProxyResp:
                handleNewProxyResponse((NewProxyRespMessage) message);
                break;
            case NewWorkConn:
                handleNewWorkConn((NewWorkConnMessage) message);
                break;
            case HeartbeatResp:
                // 心跳响应，无需处理
                break;
            case Error:
                handleError((ErrorMessage) message);
                break;
            default:
                logger.warn("未知消息类型: {}", message.getType());
        }
    }
    
    /**
     * 处理登录响应
     */
    private void handleLoginResponse(LoginRespMessage resp) {
        if (resp.isSuccess()) {
            logger.info("登录成功，运行ID: {}", resp.getRunId());
            this.runId = resp.getRunId();
            
            // 创建代理
            createProxies();
        } else {
            logger.error("登录失败: {}", resp.getError());
        }
    }
    
    /**
     * 创建代理
     */
    private void createProxies() {
        for (FrpcConfig.ProxyConfig proxyConfig : config.getProxies()) {
            try {
                createProxy(proxyConfig);
            } catch (Exception e) {
                logger.error("创建代理 {} 失败: {}", proxyConfig.getName(), e.getMessage());
            }
        }
    }
    
    /**
     * 创建单个代理
     */
    private void createProxy(FrpcConfig.ProxyConfig proxyConfig) throws IOException {
        NewProxyMessage newProxyMsg = new NewProxyMessage();
        newProxyMsg.setProxyName(proxyConfig.getName());
        newProxyMsg.setProxyType(proxyConfig.getType());
        newProxyMsg.setUseEncryption(proxyConfig.isUseEncryption());
        newProxyMsg.setUseCompression(proxyConfig.isUseCompression());
        
        if ("tcp".equals(proxyConfig.getType()) || "udp".equals(proxyConfig.getType())) {
            newProxyMsg.setLocalIp(proxyConfig.getLocalIp());
            newProxyMsg.setLocalPort(proxyConfig.getLocalPort());
            newProxyMsg.setRemotePort(proxyConfig.getRemotePort());
        } else if ("http".equals(proxyConfig.getType()) || "https".equals(proxyConfig.getType())) {
            newProxyMsg.setLocalIp(proxyConfig.getLocalIp());
            newProxyMsg.setLocalPort(proxyConfig.getLocalPort());
            newProxyMsg.setCustomDomains(proxyConfig.getCustomDomains());
        }
        
        sendMessage(newProxyMsg);
    }
    
    /**
     * 处理新建代理响应
     */
    private void handleNewProxyResponse(NewProxyRespMessage resp) {
        if (resp.isSuccess()) {
            logger.info("代理 {} 创建成功，远程地址: {}:{}", 
                resp.getProxyName(), resp.getRemoteAddr(), resp.getRemotePort());
            
            // 启动代理处理器
            FrpcConfig.ProxyConfig proxyConfig = findProxyConfig(resp.getProxyName());
            if (proxyConfig != null) {
                ProxyHandler handler = new ProxyHandler(proxyConfig, this);
                proxyHandlers.put(resp.getProxyName(), handler);
                try {
                    handler.start();
                } catch (Exception e) {
                    logger.error("启动代理处理器失败", e);
                }
            }
        } else {
            logger.error("代理 {} 创建失败: {}", resp.getProxyName(), resp.getError());
        }
    }
    
    /**
     * 处理新建工作连接请求
     */
    private void handleNewWorkConn(NewWorkConnMessage message) {
        ProxyHandler handler = proxyHandlers.get(message.getProxyName());
        if (handler != null) {
            try {
                handler.createWorkConnection(message);
                
                // 发送响应
                NewWorkConnRespMessage resp = new NewWorkConnRespMessage(
                    message.getProxyName(), 
                    null
                );
                sendMessage(resp);
            } catch (Exception e) {
                logger.error("创建工作连接失败", e);
                
                NewWorkConnRespMessage resp = new NewWorkConnRespMessage(
                    message.getProxyName(), 
                    e.getMessage()
                );
                try {
                    sendMessage(resp);
                } catch (IOException ex) {
                    logger.error("发送错误响应失败", ex);
                }
            }
        }
    }
    
    /**
     * 处理错误消息
     */
    private void handleError(ErrorMessage error) {
        logger.error("收到错误消息: {}", error.getError());
    }
    
    /**
     * 发送消息
     */
    public void sendMessage(Message message) throws IOException {
        byte[] data = MessageCodec.encode(message);
        ByteBuffer buffer = ByteBuffer.wrap(data);
        
        synchronized (controlChannel) {
            controlChannel.write(buffer);
        }
    }
    
    /**
     * 重连
     */
    private void reconnect() throws IOException {
        logger.info("正在重连...");
        stop();
        
        // 清理代理处理器
        for (ProxyHandler handler : proxyHandlers.values()) {
            handler.stop();
        }
        proxyHandlers.clear();
        
        // 重新连接
        connect();
    }
    
    /**
     * 查找代理配置
     */
    private FrpcConfig.ProxyConfig findProxyConfig(String name) {
        for (FrpcConfig.ProxyConfig proxy : config.getProxies()) {
            if (proxy.getName().equals(name)) {
                return proxy;
            }
        }
        return null;
    }
    
    /**
     * 获取运行ID
     */
    public String getRunId() {
        return runId;
    }
    
    /**
     * 检查是否正在运行
     */
    public boolean isRunning() {
        return running.get();
    }
    
    /**
     * 获取配置
     */
    public FrpcConfig getConfig() {
        return config;
    }
}